package com.packer.walle;

import java.io.DataOutput;
import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.FileChannel;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Set;

public class ApkV2Utils{

    private ApkV2Utils(){
    }

    private static long findZipCommentLength(FileChannel fileChannel) throws IOException{
        long archiveSize = fileChannel.size();
        if(archiveSize < 22L){
            throw new IOException("APK too small for ZIP End of Central Directory (EOCD) record");
        }else{
            long maxCommentLength = Math.min(archiveSize - 22L, 65535L);
            long eocdWithEmptyCommentStartPosition = archiveSize - 22L;

            for(int expectedCommentLength = 0; (long) expectedCommentLength <= maxCommentLength; ++expectedCommentLength){
                long eocdStartPos = eocdWithEmptyCommentStartPosition - (long) expectedCommentLength;
                ByteBuffer byteBuffer = ByteBuffer.allocate(4);
                fileChannel.position(eocdStartPos);
                fileChannel.read(byteBuffer);
                byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
                if(byteBuffer.getInt(0) == 101010256){
                    ByteBuffer commentLengthByteBuffer = ByteBuffer.allocate(2);
                    fileChannel.position(eocdStartPos + 20L);
                    fileChannel.read(commentLengthByteBuffer);
                    commentLengthByteBuffer.order(ByteOrder.LITTLE_ENDIAN);
                    int actualCommentLength = commentLengthByteBuffer.getShort(0);
                    if(actualCommentLength == expectedCommentLength){
                        return actualCommentLength;
                    }
                }
            }
            throw new IOException("ZIP End of Central Directory (EOCD) record not found");
        }
    }

    private static long findCentralDirStartOffset(FileChannel fileChannel) throws IOException{
        return findCentralDirStartOffset(fileChannel, findZipCommentLength(fileChannel));
    }

    private static long findCentralDirStartOffset(FileChannel fileChannel, long commentLength) throws IOException{
        ByteBuffer zipCentralDirectoryStart = ByteBuffer.allocate(4);
        zipCentralDirectoryStart.order(ByteOrder.LITTLE_ENDIAN);
        fileChannel.position(fileChannel.size() - commentLength - 6L);
        fileChannel.read(zipCentralDirectoryStart);
        return zipCentralDirectoryStart.getInt(0);
    }

    private static Pair<ByteBuffer, Long> findApkSigningBlock(FileChannel fileChannel) throws IOException{
        long centralDirOffset = findCentralDirStartOffset(fileChannel);
        return findApkSigningBlock(fileChannel, centralDirOffset);
    }

    private static Pair<ByteBuffer, Long> findApkSigningBlock(FileChannel fileChannel, long centralDirOffset) throws IOException{
        if(centralDirOffset < 32L){
            throw new V2SignException("APK too small for APK Signing Block. ZIP Central Directory offset: " + centralDirOffset);
        }else{
            fileChannel.position(centralDirOffset - 24L);
            ByteBuffer footer = ByteBuffer.allocate(24);
            fileChannel.read(footer);
            footer.order(ByteOrder.LITTLE_ENDIAN);
            if(footer.getLong(8) == 2334950737559900225L && footer.getLong(16) == 3617552046287187010L){
                long apkSigBlockSizeInFooter = footer.getLong(0);
                if(apkSigBlockSizeInFooter >= (long) footer.capacity() && apkSigBlockSizeInFooter <= 2147483639L){
                    int totalSize = (int) (apkSigBlockSizeInFooter + 8L);
                    long apkSigBlockOffset = centralDirOffset - (long) totalSize;
                    if(apkSigBlockOffset < 0L){
                        throw new V2SignException("APK Signing Block offset out of range: " + apkSigBlockOffset);
                    }else{
                        fileChannel.position(apkSigBlockOffset);
                        ByteBuffer apkSigBlock = ByteBuffer.allocate(totalSize);
                        fileChannel.read(apkSigBlock);
                        apkSigBlock.order(ByteOrder.LITTLE_ENDIAN);
                        long apkSigBlockSizeInHeader = apkSigBlock.getLong(0);
                        if(apkSigBlockSizeInHeader != apkSigBlockSizeInFooter){
                            throw new V2SignException("APK Signing Block sizes in header and footer do not match: " + apkSigBlockSizeInHeader + " vs " + apkSigBlockSizeInFooter);
                        }else{
                            return Pair.of(apkSigBlock, apkSigBlockOffset);
                        }
                    }
                }else{
                    throw new V2SignException("APK Signing Block size out of range: " + apkSigBlockSizeInFooter);
                }
            }else{
                throw new V2SignException("No APK Signing Block before ZIP Central Directory.");
            }
        }
    }

    private static LinkedHashMap<Integer, ByteBuffer> findIdValues(ByteBuffer apkSigningBlock) throws IOException{
        checkByteOrderLittleEndian(apkSigningBlock);
        ByteBuffer pairs = sliceFromTo(apkSigningBlock, 8, apkSigningBlock.capacity() - 24);
        LinkedHashMap<Integer, ByteBuffer> idValues = new LinkedHashMap<>();
        int entryCount = 0;

        while(pairs.hasRemaining()){
            ++entryCount;
            if(pairs.remaining() < 8){
                throw new IOException("Insufficient data to read size of APK Signing Block entry #" + entryCount);
            }

            long lenLong = pairs.getLong();
            if(lenLong < 4L || lenLong > 2147483647L){
                throw new IOException("APK Signing Block entry #" + entryCount + " size out of range: " + lenLong);
            }

            int len = (int) lenLong;
            int nextEntryPos = pairs.position() + len;
            if(len > pairs.remaining()){
                throw new IOException("APK Signing Block entry #" + entryCount + " size out of range: " + len + ", available: " + pairs.remaining());
            }

            int id = pairs.getInt();
            idValues.put(id, getByteBuffer(pairs, len - 4));
            pairs.position(nextEntryPos);
        }
        return idValues;
    }

    private static ByteBuffer sliceFromTo(ByteBuffer source, int start, int end){
        if(start < 0){
            throw new IllegalArgumentException("start: " + start);
        }else if(end < start){
            throw new IllegalArgumentException("end < start: " + end + " < " + start);
        }else{
            int capacity = source.capacity();
            if(end > source.capacity()){
                throw new IllegalArgumentException("end > capacity: " + end + " > " + capacity);
            }else{
                int originalLimit = source.limit();
                int originalPosition = source.position();

                ByteBuffer var7;
                try{
                    source.position(0);
                    source.limit(end);
                    source.position(start);
                    ByteBuffer result = source.slice();
                    result.order(source.order());
                    var7 = result;
                }finally{
                    source.position(0);
                    source.limit(originalLimit);
                    source.position(originalPosition);
                }
                return var7;
            }
        }
    }

    private static ByteBuffer getByteBuffer(ByteBuffer source, int size) throws IOException{
        if(size < 0){
            throw new IOException("size: " + size);
        }else{
            int originalLimit = source.limit();
            int position = source.position();
            int limit = position + size;
            if(limit >= position && limit <= originalLimit){
                source.limit(limit);

                ByteBuffer var6;
                try{
                    ByteBuffer result = source.slice();
                    result.order(source.order());
                    source.position(limit);
                    var6 = result;
                }finally{
                    source.limit(originalLimit);
                }
                return var6;
            }else{
                throw new IOException();
            }
        }
    }

    private static void checkByteOrderLittleEndian(ByteBuffer buffer) throws IOException{
        try{
            if(buffer.order() != ByteOrder.LITTLE_ENDIAN){
                throw new IOException("ByteBuffer byte order must be little endian");
            }
        }catch(Exception e){
            throw new IOException("ByteBuffer byte order must be little endian.");
        }
    }

    public static String byteBufferToString(java.nio.ByteBuffer buffer){
        byte[] bytes = byteBufferToBytes(buffer);
        if(bytes.length > 0){
            return new String(bytes);
        }
        return "";
    }

    public static byte[] byteBufferToBytes(ByteBuffer buf){
        if(buf == null){
            return new byte[]{};
        }
        byte[] array = buf.array();
        int arrayOffset = buf.arrayOffset();
        return Arrays.copyOfRange(array, arrayOffset + buf.position(), arrayOffset + buf.limit());
    }

    public static ByteBuffer bytesToByteBuffer(byte[] bytes){
        if(bytes == null){
            bytes = new byte[]{};
        }
        ByteBuffer byteBuffer = ByteBuffer.allocate(bytes.length);
        byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
        byteBuffer.put(bytes, 0, bytes.length);
        byteBuffer.flip();
        return byteBuffer;
    }

    public static void writeApkSigningBlock(File apkFile, ApkSigningBlockHandler handler) throws IOException{

        try(RandomAccessFile raf = new RandomAccessFile(apkFile, "rw");
            FileChannel fc = raf.getChannel();){

            long commentLength = ApkV2Utils.findZipCommentLength(fc);
            long centralDirStartOffset = ApkV2Utils.findCentralDirStartOffset(fc, commentLength);
            Pair<ByteBuffer, Long> apkSigningBlockAndOffset = ApkV2Utils.findApkSigningBlock(fc, centralDirStartOffset);
            ByteBuffer apkSigningBlock2 = apkSigningBlockAndOffset.getFirst();
            long apkSigningBlockOffset = apkSigningBlockAndOffset.getSecond();
            if(centralDirStartOffset == 0L || apkSigningBlockOffset == 0L){
                throw new IOException("No APK Signature Scheme v2 block in APK Signing Block");
            }

            LinkedHashMap<Integer, ByteBuffer> originIdValues = ApkV2Utils.findIdValues(apkSigningBlock2);
            ByteBuffer apkSignatureSchemeV2Block = originIdValues.get(1896449818);
            if(apkSignatureSchemeV2Block == null){
                throw new IOException("No APK Signature Scheme v2 block in APK Signing Block");
            }

            LinkedHashMap<Integer, ByteBuffer> newBlocks = handler.handleBlocks(originIdValues);

            raf.seek(centralDirStartOffset);
            byte[] centralDirBytes = new byte[(int) (fc.size() - centralDirStartOffset)];
            raf.read(centralDirBytes);
            fc.position(apkSigningBlockOffset);
            long length = writeTo(raf, newBlocks);
            raf.write(centralDirBytes);
            raf.setLength(raf.getFilePointer());
            raf.seek(fc.size() - commentLength - 6L);
            ByteBuffer temp = ByteBuffer.allocate(4);
            temp.order(ByteOrder.LITTLE_ENDIAN);
            temp.putInt((int) (centralDirStartOffset + length + 8L - (centralDirStartOffset - apkSigningBlockOffset)));
            temp.flip();
            raf.write(temp.array());
        }
    }

    public static ByteBuffer readBlockById(File apkFile, int blockId){
        LinkedHashMap<Integer, ByteBuffer> blockList = readBlockMap(apkFile);
        return blockList == null ? null : blockList.get(blockId);
    }

    public static LinkedHashMap<Integer, ByteBuffer> readBlockMap(File apkFile){

        LinkedHashMap<Integer, ByteBuffer> blocks;

        try(RandomAccessFile raf = new RandomAccessFile(apkFile, "r"); FileChannel fc = raf.getChannel()){
            ByteBuffer apkSigningBlock = ApkV2Utils.findApkSigningBlock(fc).getFirst();
            blocks = ApkV2Utils.findIdValues(apkSigningBlock);
        }catch(IOException e){
            blocks = null;
        }
        return blocks;
    }

    private static long writeTo(DataOutput dataOutput, HashMap<Integer, ByteBuffer> map) throws IOException{
        long length = 24L;

        Set<Integer> keySet = map.keySet();
        for(int key : keySet){
            byte[] bytes = ApkV2Utils.byteBufferToBytes(map.get(key));
            length += 12 + bytes.length;
        }

        ByteBuffer byteBuffer = ByteBuffer.allocate(8);
        byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
        byteBuffer.putLong(length);
        byteBuffer.flip();
        dataOutput.write(byteBuffer.array());

        for(int key : keySet){
            byte[] bytes = ApkV2Utils.byteBufferToBytes(map.get(key));
            byteBuffer = ByteBuffer.allocate(8);
            byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
            byteBuffer.putLong(bytes.length + 4);
            byteBuffer.flip();
            dataOutput.write(byteBuffer.array());
            byteBuffer = ByteBuffer.allocate(4);
            byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
            byteBuffer.putInt(key);
            byteBuffer.flip();
            dataOutput.write(byteBuffer.array());
            dataOutput.write(bytes);
        }

        byteBuffer = ByteBuffer.allocate(8);
        byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
        byteBuffer.putLong(length);
        byteBuffer.flip();
        dataOutput.write(byteBuffer.array());
        byteBuffer = ByteBuffer.allocate(8);
        byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
        byteBuffer.putLong(2334950737559900225L);
        byteBuffer.flip();
        dataOutput.write(byteBuffer.array());
        byteBuffer = ByteBuffer.allocate(8);
        byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
        byteBuffer.putLong(3617552046287187010L);
        byteBuffer.flip();
        dataOutput.write(byteBuffer.array());
        return length;
    }

    public interface ApkSigningBlockHandler{
        LinkedHashMap<Integer, ByteBuffer> handleBlocks(LinkedHashMap<Integer, ByteBuffer> originIdValues);
    }
}